import numpy as np
from PIL import Image


def ideal_lowpass_filter(image_path, output_path, cutoff_frequency):
    """
    对输入的RGB图像应用理想低通滤波器。

    参数:
    image_path (str): 输入图像的文件路径。
    output_path (str): 输出滤波后图像的文件路径。
    cutoff_frequency (float): 截止频率，用于确定滤波器的通带范围，单位通常根据图像尺寸等因素相关，这里基于图像的归一化频率。
    """
    # 1. 打开图像并转换为numpy数组
    image = Image.open(image_path)
    width, height = image.size
    image_array = np.array(image)

    #  对每个颜色通道分别进行处理（RGB图像有3个通道）
    filtered_image_array = np.zeros_like(image_array)
    for channel in range(3):
        # 获取当前通道的二维数组
        channel_data = image_array[:, :, channel]

        # 进行二维离散傅里叶变换（DFT），将空域转换到频域
        dft = np.fft.fft2(channel_data)

        # 将零频率分量移到频谱中心，便于后续处理和可视化
        dft_shifted = np.fft.fftshift(dft)

        # 获取频谱的行数和列数
        rows, cols = dft_shifted.shape

        # 计算频域中每个点到频谱中心的距离（用于判断是否在通带内）
        crow, ccol = rows // 2, cols // 2
        distance_matrix = np.sqrt((np.arange(rows)[:, np.newaxis] - crow) ** 2 + (np.arange(cols) - ccol) ** 2)

        # 根据截止频率构建理想低通滤波器的掩模（在通带内为1，通带外为0）
        mask = np.where(distance_matrix <= cutoff_frequency, 1, 0)

        # 将滤波器掩模应用到频域数据上（频域相乘相当于空域卷积）
        filtered_dft_shifted = dft_shifted * mask

        # 将零频率分量移回原来的位置
        filtered_dft = np.fft.ifftshift(filtered_dft_shifted)

        # 进行二维离散傅里叶逆变换（IDFT），将频域转换回空域
        filtered_channel_data = np.fft.ifft2(filtered_dft)

        # 取实部并转换为合适的数据类型（uint8）
        filtered_channel_data = np.real(filtered_channel_data).astype(np.uint8)

        # 将滤波后的通道数据放回原图像数组对应的位置
        filtered_image_array[:, :, channel] = filtered_channel_data

    # 3. 将滤波后的numpy数组转换回图像并保存
    filtered_image = Image.fromarray(filtered_image_array)
    filtered_image.save(output_path)
    return filtered_image

image_path = "D:/cangku/computer-image-project-design/test.jpg"
output_path = "D:/cangku/computer-image-project-design/test1.jpg"
cutoff_frequency = 100
ideal_lowpass_filter(image_path, output_path, cutoff_frequency)